from typing import Dict, List, Optional, Any
import asyncio
import aiohttp
import numpy as np
from .data_storage import DataStorage

class CommunicationManager:
    """Manages network communication for worker nodes.
    
    This class handles all network communications including data transfer
    with master node and other workers, task status updates, and data
    synchronization.
    """
    
    def __init__(self, worker_id: str, master_url: str,
                 data_storage: DataStorage, config: Dict[str, Any]):
        """Initialize the communication manager.
        
        Args:
            worker_id: Unique identifier for this worker
            master_url: URL of the master node
            data_storage: Data storage manager for handling cached data
            config: Configuration parameters
        """
        self.worker_id = worker_id
        self.master_url = master_url
        self.data_storage = data_storage
        self.config = config
        
        self.timeout = config['network']['timeout']
        self.retry_attempts = config['network']['retry_attempts']
        self.buffer_size = config['network']['buffer_size']
        
        self.session: Optional[aiohttp.ClientSession] = None
    
    async def start(self) -> None:
        """Start the communication manager."""
        self.session = aiohttp.ClientSession()
        await self._register_with_master()
    
    async def stop(self) -> None:
        """Stop the communication manager."""
        if self.session:
            await self.session.close()
            self.session = None
    
    async def _register_with_master(self) -> None:
        """Register this worker with the master node."""
        try:
            async with self.session.post(
                f"{self.master_url}/register",
                json={
                    'worker_id': self.worker_id,
                    'capabilities': {
                        'gpu_memory': self.config.get('gpu_memory', 0),
                        'num_cores': self.config.get('num_cores', 1)
                    }
                },
                timeout=self.timeout
            ) as response:
                if response.status != 200:
                    raise RuntimeError(f"Failed to register with master: {response.status}")
        except Exception as e:
            print(f"Error registering with master: {str(e)}")
            raise
    
    async def get_task(self) -> Optional[Dict[str, Any]]:
        """Get a new task from the master node.
        
        Returns:
            Optional[Dict[str, Any]]: Task information if available
        """
        try:
            async with self.session.get(
                f"{self.master_url}/get_task",
                params={'worker_id': self.worker_id},
                timeout=self.timeout
            ) as response:
                if response.status == 200:
                    return await response.json()
                elif response.status == 204:
                    return None  # No tasks available
                else:
                    raise RuntimeError(f"Failed to get task: {response.status}")
        except Exception as e:
            print(f"Error getting task: {str(e)}")
            return None
    
    async def update_task_status(self, task_id: str, status: str,
                                result: Optional[Dict[str, Any]] = None) -> bool:
        """Update the status of a task with the master.
        
        Args:
            task_id: Task identifier
            status: New status of the task
            result: Optional result data
            
        Returns:
            bool: True if update was successful
        """
        try:
            async with self.session.post(
                f"{self.master_url}/update_task",
                json={
                    'worker_id': self.worker_id,
                    'task_id': task_id,
                    'status': status,
                    'result': result
                },
                timeout=self.timeout
            ) as response:
                return response.status == 200
        except Exception as e:
            print(f"Error updating task status: {str(e)}")
            return False
    
    async def fetch_data(self, data_id: str, source_worker: str) -> bool:
        """Fetch data from another worker.
        
        Args:
            data_id: Identifier of the data to fetch
            source_worker: Worker to fetch from
            
        Returns:
            bool: True if fetch was successful
        """
        try:
            # First check if we already have the data
            if self.data_storage.load_tensor(data_id) is not None:
                return True
            
            # Fetch data from source worker
            async with self.session.get(
                f"http://{source_worker}/data/{data_id}",
                timeout=self.timeout
            ) as response:
                if response.status == 200:
                    # Read data in chunks
                    data = await response.read()
                    tensor = np.frombuffer(data, dtype=np.float32)
                    
                    # Get metadata
                    metadata = response.headers.get('X-Tensor-Metadata')
                    if metadata:
                        metadata = eval(metadata)  # Convert string to dict
                        tensor = tensor.reshape(metadata['shape'])
                    
                    # Store in local cache
                    return self.data_storage.store_tensor(data_id, tensor, metadata)
                else:
                    return False
        except Exception as e:
            print(f"Error fetching data: {str(e)}")
            return False
    
    async def send_data(self, data_id: str, target_worker: str) -> bool:
        """Send data to another worker.
        
        Args:
            data_id: Identifier of the data to send
            target_worker: Worker to send to
            
        Returns:
            bool: True if send was successful
        """
        try:
            # Load data from storage
            tensor = self.data_storage.load_tensor(data_id)
            if tensor is None:
                return False
            
            metadata = self.data_storage.get_metadata(data_id)
            
            # Send data
            headers = {
                'Content-Type': 'application/octet-stream',
                'X-Tensor-Metadata': str(metadata) if metadata else ''
            }
            
            async with self.session.post(
                f"http://{target_worker}/data",
                data=tensor.tobytes(),
                headers=headers,
                timeout=self.timeout
            ) as response:
                return response.status == 200
        except Exception as e:
            print(f"Error sending data: {str(e)}")
            return False
    
    async def report_metrics(self, metrics: Dict[str, Any]) -> None:
        """Report metrics to the master node.
        
        Args:
            metrics: Dictionary of metrics to report
        """
        try:
            async with self.session.post(
                f"{self.master_url}/report_metrics",
                json={
                    'worker_id': self.worker_id,
                    'metrics': metrics
                },
                timeout=self.timeout
            ) as response:
                if response.status != 200:
                    print(f"Failed to report metrics: {response.status}")
        except Exception as e:
            print(f"Error reporting metrics: {str(e)}")
    
    async def sync_with_workers(self, sync_point: int,
                               worker_group: List[str]) -> bool:
        """Synchronize with other workers at a sync point.
        
        Args:
            sync_point: Synchronization point number
            worker_group: List of workers to synchronize with
            
        Returns:
            bool: True if synchronization was successful
        """
        try:
            # Notify master we're ready to sync
            async with self.session.post(
                f"{self.master_url}/sync_ready",
                json={
                    'worker_id': self.worker_id,
                    'sync_point': sync_point,
                    'worker_group': worker_group
                },
                timeout=self.timeout
            ) as response:
                if response.status != 200:
                    return False
                
                # Wait for all workers to be ready
                sync_data = await response.json()
                if not sync_data.get('success', False):
                    return False
                
                return True
        except Exception as e:
            print(f"Error during synchronization: {str(e)}")
            return False